{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook written by [Zhedong Zheng](https://github.com/zhedongzheng)\n",
    "\n",
    "![title](img/birnn.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import sklearn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "VOCAB_SIZE = 20000\n",
    "EMBED_DIM = 100\n",
    "RNN_SIZE = 70\n",
    "CLIP_NORM = 5.0\n",
    "BATCH_SIZE = 32\n",
    "LR = {'start': 5e-3, 'end': 5e-4, 'steps': 1500}\n",
    "N_EPOCH = 2\n",
    "N_CLASS = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sort_by_len(x, y):\n",
    "    idx = sorted(range(len(x)), key=lambda i: len(x[i]))\n",
    "    return x[idx], y[idx]\n",
    "\n",
    "def pad_sentence_batch(sent_batch, thres=400):\n",
    "    max_seq_len = max([len(sent) for sent in sent_batch])\n",
    "    if max_seq_len > thres:\n",
    "        max_seq_len = thres\n",
    "        sent_batch = [sent[-thres:] for sent in sent_batch]\n",
    "    padded_seqs = [(sent + [0]*(max_seq_len - len(sent))) for sent in sent_batch]\n",
    "    return padded_seqs\n",
    "\n",
    "def next_train_batch(X_train, y_train):\n",
    "    for i in range(0, len(X_train), BATCH_SIZE):\n",
    "        padded_seqs = pad_sentence_batch(X_train[i : i+BATCH_SIZE])\n",
    "        yield padded_seqs, y_train[i : i+BATCH_SIZE]\n",
    "        \n",
    "def next_test_batch(X_test):\n",
    "    for i in range(0, len(X_test), BATCH_SIZE):\n",
    "        padded_seqs = pad_sentence_batch(X_test[i : i+BATCH_SIZE])\n",
    "        yield padded_seqs\n",
    "        \n",
    "def train_input_fn(X_train, y_train):\n",
    "    dataset = tf.data.Dataset.from_generator(\n",
    "        lambda: next_train_batch(X_train, y_train),\n",
    "        (tf.int32, tf.int64),\n",
    "        (tf.TensorShape([None,None]), tf.TensorShape([None])))\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    return iterator.get_next()\n",
    "\n",
    "def predict_input_fn(X_test):\n",
    "    dataset = tf.data.Dataset.from_generator(\n",
    "        lambda: next_test_batch(X_test),\n",
    "        tf.int32,\n",
    "        tf.TensorShape([None,None]))\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    return iterator.get_next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rnn_cell():\n",
    "    return tf.nn.rnn_cell.GRUCell(RNN_SIZE, kernel_initializer=tf.orthogonal_initializer())\n",
    "\n",
    "def forward(inputs, mode):\n",
    "    is_training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
    "    \n",
    "    x = tf.contrib.layers.embed_sequence(inputs, VOCAB_SIZE, EMBED_DIM)\n",
    "    x = tf.layers.dropout(x, 0.2, training=is_training)\n",
    "\n",
    "    bi_outs, bi_states = tf.nn.bidirectional_dynamic_rnn(\n",
    "        rnn_cell(), rnn_cell(), x, tf.count_nonzero(inputs, 1), dtype=tf.float32)\n",
    "    \n",
    "    x = tf.concat(bi_states, -1)\n",
    "    logits = tf.layers.dense(x, N_CLASS)\n",
    "    return logits\n",
    "\n",
    "def clip_grads(loss):\n",
    "    params = tf.trainable_variables()\n",
    "    grads = tf.gradients(loss, params)\n",
    "    clipped_grads, _ = tf.clip_by_global_norm(grads, CLIP_NORM)\n",
    "    return zip(clipped_grads, params)\n",
    "\n",
    "def model_fn(features, labels, mode):\n",
    "    logits = forward(features, mode)\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        preds = tf.argmax(logits, -1)\n",
    "        return tf.estimator.EstimatorSpec(mode, predictions=preds)\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        global_step = tf.train.get_global_step()\n",
    "\n",
    "        lr_op = tf.train.exponential_decay(\n",
    "            LR['start'], global_step, LR['steps'], LR['end']/LR['start'])\n",
    "\n",
    "        loss_op = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "            logits=logits, labels=labels))\n",
    "\n",
    "        train_op = tf.train.AdamOptimizer(lr_op).apply_gradients(\n",
    "            clip_grads(loss_op), global_step=global_step)\n",
    "\n",
    "        lth = tf.train.LoggingTensorHook({'lr': lr_op}, every_n_iter=100)\n",
    "        \n",
    "        return tf.estimator.EstimatorSpec(\n",
    "            mode=mode, loss=loss_op, train_op=train_op, training_hooks=[lth])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using default config.\n",
      "WARNING:tensorflow:Using temporary folder as model directory: /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp0halwsrb\n",
      "INFO:tensorflow:Using config: {'_model_dir': '/var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp0halwsrb', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': None, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x1132385f8>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 1 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp0halwsrb/model.ckpt.\n",
      "INFO:tensorflow:loss = 0.6929605, step = 1\n",
      "INFO:tensorflow:lr = 0.005\n",
      "INFO:tensorflow:global_step/sec: 15.9543\n",
      "INFO:tensorflow:loss = 0.2603664, step = 101 (6.269 sec)\n",
      "INFO:tensorflow:lr = 0.004288479 (6.269 sec)\n",
      "INFO:tensorflow:global_step/sec: 10.7623\n",
      "INFO:tensorflow:loss = 0.4148208, step = 201 (9.292 sec)\n",
      "INFO:tensorflow:lr = 0.0036782112 (9.292 sec)\n",
      "INFO:tensorflow:global_step/sec: 9.36032\n",
      "INFO:tensorflow:loss = 0.32667416, step = 301 (10.683 sec)\n",
      "INFO:tensorflow:lr = 0.0031547868 (10.683 sec)\n",
      "INFO:tensorflow:global_step/sec: 7.66704\n",
      "INFO:tensorflow:loss = 0.4845091, step = 401 (13.043 sec)\n",
      "INFO:tensorflow:lr = 0.0027058476 (13.043 sec)\n",
      "INFO:tensorflow:global_step/sec: 6.10341\n",
      "INFO:tensorflow:loss = 0.32211247, step = 501 (16.384 sec)\n",
      "INFO:tensorflow:lr = 0.0023207942 (16.385 sec)\n",
      "INFO:tensorflow:global_step/sec: 4.85513\n",
      "INFO:tensorflow:loss = 0.22540471, step = 601 (20.597 sec)\n",
      "INFO:tensorflow:lr = 0.0019905358 (20.597 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.62032\n",
      "INFO:tensorflow:loss = 0.31916237, step = 701 (27.622 sec)\n",
      "INFO:tensorflow:lr = 0.0017072745 (27.622 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 782 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp0halwsrb/model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 0.0733175.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp0halwsrb/model.ckpt-782\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "\n",
      "Validation Accuracy: 0.9048\n",
      "\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp0halwsrb/model.ckpt-782\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 783 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp0halwsrb/model.ckpt.\n",
      "INFO:tensorflow:loss = 0.31547743, step = 783\n",
      "INFO:tensorflow:lr = 0.0015053472\n",
      "INFO:tensorflow:global_step/sec: 15.5694\n",
      "INFO:tensorflow:loss = 0.10211637, step = 883 (6.424 sec)\n",
      "INFO:tensorflow:lr = 0.00129113 (6.424 sec)\n",
      "INFO:tensorflow:global_step/sec: 10.9138\n",
      "INFO:tensorflow:loss = 0.057801604, step = 983 (9.163 sec)\n",
      "INFO:tensorflow:lr = 0.001107397 (9.163 sec)\n",
      "INFO:tensorflow:global_step/sec: 9.35336\n",
      "INFO:tensorflow:loss = 0.058652263, step = 1083 (10.691 sec)\n",
      "INFO:tensorflow:lr = 0.00094980985 (10.691 sec)\n",
      "INFO:tensorflow:global_step/sec: 8.21201\n",
      "INFO:tensorflow:loss = 0.107295245, step = 1183 (12.178 sec)\n",
      "INFO:tensorflow:lr = 0.00081464805 (12.178 sec)\n",
      "INFO:tensorflow:global_step/sec: 6.66586\n",
      "INFO:tensorflow:loss = 0.05238443, step = 1283 (15.001 sec)\n",
      "INFO:tensorflow:lr = 0.0006987203 (15.001 sec)\n",
      "INFO:tensorflow:global_step/sec: 5.00387\n",
      "INFO:tensorflow:loss = 0.0082637975, step = 1383 (19.985 sec)\n",
      "INFO:tensorflow:lr = 0.00059928955 (19.985 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.67222\n",
      "INFO:tensorflow:loss = 0.08863423, step = 1483 (27.231 sec)\n",
      "INFO:tensorflow:lr = 0.00051400816 (27.231 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 1564 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp0halwsrb/model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 0.0010443431.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp0halwsrb/model.ckpt-1564\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "\n",
      "Validation Accuracy: 0.8983\n",
      "\n"
     ]
    }
   ],
   "source": [
    "(X_train, y_train), (X_test, y_test) = tf.keras.datasets.imdb.load_data(num_words=VOCAB_SIZE)\n",
    "X_train, y_train = sort_by_len(X_train, y_train)\n",
    "X_test, y_test = sort_by_len(X_test, y_test)\n",
    "\n",
    "estimator = tf.estimator.Estimator(model_fn)\n",
    "\n",
    "for _ in range(N_EPOCH):\n",
    "    estimator.train(lambda: train_input_fn(X_train, y_train))\n",
    "    y_pred = np.fromiter(estimator.predict(lambda: predict_input_fn(X_test)), np.int32)\n",
    "    print(\"\\nValidation Accuracy: %.4f\\n\" % (y_pred==y_test).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
